//
//  ZYSecurityPolicy.m
//  ZYWebImage
//
//  Created by wangzhipan on 2025/4/11.
//

#import "ZYSecurityPolicy.h"

@interface ZYSecurityPolicy()
@property (readwrite, nonatomic, strong) NSSet *pinnedPublicKeys;
@end

@implementation ZYSecurityPolicy
//返回包含在bundle中的任何证书
+ (NSSet <NSData *> *)certificatesInBundle:(NSBundle *)bundle
{
    NSArray *paths = [bundle pathsForResourcesOfType:@".cer" inDirectory:@"."];
    NSMutableSet *cerSet = [NSMutableSet setWithCapacity:paths.count];
    
    for (NSString *path in paths) {
        NSData *data = [NSData dataWithContentsOfFile:path];
        [cerSet addObject:data];
    }
    return cerSet;
}

+ (instancetype)defaultPolicy
{
    ZYSecurityPolicy *policy = [[ZYSecurityPolicy alloc] init];
    policy.SSLPinningMode = ZYSSLPinningModeNone;
    return policy;
}

+ (instancetype)policyWithSSLPinningMode:(ZYSSLPinningMode)SSLPinningMode
{
    NSSet <NSData *> *cerSet = [self certificatesInBundle:[NSBundle mainBundle]];
    return [self policyWithSSLPinningMode:SSLPinningMode
                   withPinnedCertificates:cerSet];
}

+ (instancetype)policyWithSSLPinningMode:(ZYSSLPinningMode)SSLPinningMode
                  withPinnedCertificates:(NSSet <NSData *> *)pinnedCertificates
{
    ZYSecurityPolicy *policy = [[ZYSecurityPolicy alloc] init];
    policy.SSLPinningMode = SSLPinningMode;
    policy.pinnedCertificates = pinnedCertificates;
    return policy;
}

- (instancetype)init
{
    if (self = [super init]) {
        self.validatesDomainName = YES;
    }
    return self;
}

- (BOOL)evaluateServerTrust:(SecTrustRef)serverTrust
                  forDomain:(nullable NSString *)domain
{
    if (domain && self.allowInvalidCertificates && self.validatesDomainName && (self.SSLPinningMode == ZYSSLPinningModeNone || [self.pinnedCertificates count] == 0)) {
        NSLog(@"In order to validate a domain name for self signed certificates, you MUST use pinning.");
        return NO;
    }
    
    NSMutableArray *policies = [NSMutableArray array];
    if (self.validatesDomainName) {
        [policies addObject:(__bridge_transfer id)SecPolicyCreateSSL(true, (__bridge CFStringRef)domain)];
    }
    else {
        [policies addObject:(__bridge_transfer id)SecPolicyCreateBasicX509()];
    }
    
    SecTrustSetPolicies(serverTrust, (__bridge CFTypeRef _Nonnull)(policies));
    
    if (self.SSLPinningMode == ZYSSLPinningModeNone) {
        //当没有设置验证模式时，检测下服务器证书是否有效，或者客户端允许无效证书
        return self.allowInvalidCertificates || [self p_serverTrustIsValid:serverTrust];
    }
    else if (!self.allowInvalidCertificates && ![self p_serverTrustIsValid:serverTrust]){
        //设置了验证模式， 如果不验证证书，但服务器证书无效
        return NO;
    }
    
    switch (self.SSLPinningMode) {
        case ZYSSLPinningModePublicKey:     //公钥验证
        {
            NSUInteger trustedPublicKeyCount = 0;
            //获取服务器所有公钥
            NSArray *publicKeys = [self p_publicKeyTrustChainForServerTrust:serverTrust];
            
            for (id trustPublicKey in publicKeys) {
                for (id pinnedKey in self.pinnedPublicKeys) {
                    //判断本地证书公钥与服务器证书的公钥是否一致
                    if ([self p_secKeyIsEqualToKey:(__bridge SecKeyRef)pinnedKey key2:(__bridge SecKeyRef)trustPublicKey]) {
                        trustedPublicKeyCount++;
                    }
                }
            }
            return trustedPublicKeyCount > 0;
        }
        case ZYSSLPinningModeCertificate:       //证书验证
        {
            NSMutableArray *pinnedCertificates = [NSMutableArray array];
            for (NSData *data in self.pinnedCertificates) {
                [pinnedCertificates addObject:(__bridge_transfer id)SecCertificateCreateWithData(NULL, (__bridge CFDataRef)data)];
            }
            
            //指定一组证书作为信任锚点，用于证书链验证过程。信任锚点通常是根证书或自签名证书，它们是证书信任链的起点
            SecTrustSetAnchorCertificates(serverTrust, (__bridge CFArrayRef)pinnedCertificates);
            
            //验证服务器证书是否有效
            if (![self p_serverTrustIsValid:serverTrust]) {
                return NO;
            }
            // 获取服务器证书链
            NSArray *serverCertificates = [self p_certificateTrustChainForServerTrust:serverTrust];
            
            //
            for (NSData *trustChainCertificate in serverCertificates.reverseObjectEnumerator) {
                //本地客户端证书包含服务器的证书
                if ([self.pinnedCertificates containsObject:trustChainCertificate]) {
                    return YES;
                }
            }
            return NO;
        }
        default:
            return NO;
    }
    return NO;
}


#pragma mark - private

// 从 NSData 格式的证书中提取公钥
- (id)p_publicKeyFromCertificateData:(NSData *)certificateData {
    id publicKey = nil;
    
    // 步骤1: 从证书数据创建 SecCertificateRef 对象
    SecCertificateRef certificate = SecCertificateCreateWithData(NULL, (__bridge CFDataRef)certificateData);
    if (certificate == NULL) {
        NSLog(@"无法从数据创建证书对象");
        return nil;
    }
    
    // 步骤2: 创建包含证书的临时 SecTrustRef 对象
    SecPolicyRef policy = SecPolicyCreateBasicX509();
    SecTrustRef trust = NULL;
    OSStatus status = SecTrustCreateWithCertificates(certificate, policy, &trust);
    
    if (status != errSecSuccess) {
        NSLog(@"创建 Trust 对象失败，错误码: %d", (int)status);
        CFRelease(certificate);
        CFRelease(policy);
        return nil;
    }
    
    // 步骤3: 评估信任对象以确保证书链完整
    SecTrustResultType trustResult;
    status = SecTrustEvaluate(trust, &trustResult);
    
    if (status != errSecSuccess) {
        NSLog(@"评估证书失败，错误码: %d", (int)status);
        CFRelease(certificate);
        CFRelease(policy);
        CFRelease(trust);
        return nil;
    }
    
    // 步骤4: 从信任对象中提取公钥
    publicKey = (__bridge_transfer id)SecTrustCopyPublicKey(trust);
    
    if (publicKey == NULL) {
        NSLog(@"无法从证书中提取公钥");
    }
    
    // 步骤5: 释放创建的对象
    CFRelease(certificate);
    CFRelease(policy);
    CFRelease(trust);
    
    return publicKey;
}

/*
 检查证书是否由受信任的根证书颁发机构(CA)签发、有没过期、有没被吊销
 */
- (BOOL)p_serverTrustIsValid:(SecTrustRef)serverTrust
{
    BOOL isValid = NO;
    SecTrustResultType result;
    
    // 评估服务器信任对象
    OSStatus status = SecTrustEvaluate(serverTrust, &result);
    
    // 检查评估是否成功
    if (status == errSecSuccess) {
        // 根据评估结果判断证书是否可信
        // kSecTrustResultUnspecified: 证书有效且被系统信任
        // kSecTrustResultProceed: 用户明确信任了这个证书
        isValid = (result == kSecTrustResultUnspecified || result == kSecTrustResultProceed);
    }
    
    return isValid;
}

/*
  获取服务器证书链
 
 从 serverTrust 对象中获取完整的证书链，并将每个证书转换为 NSData 对象。这在实现证书固定(Certificate Pinning)功能时非常重要，因为它允许开发者获取并验证服务器提供的整个证书链
 */
- (NSArray *)p_certificateTrustChainForServerTrust:(SecTrustRef)serverTrust
{
    // 获取证书链中的证书数量
    CFIndex certificateCount = SecTrustGetCertificateCount(serverTrust);
    
    // 创建可变数组存储证书数据
    NSMutableArray *trustChain = [NSMutableArray arrayWithCapacity:(NSUInteger)certificateCount];
    
    // 遍历证书链中的每个证书
    for (CFIndex i = 0; i < certificateCount; i++) {
        // 获取当前索引的证书
        SecCertificateRef certificate = SecTrustGetCertificateAtIndex(serverTrust, i);
        
        // 将证书转换为NSData格式
        NSData *certificateData = (__bridge_transfer NSData *)SecCertificateCopyData(certificate);
        
        // 将证书数据添加到数组
        [trustChain addObject:certificateData];
    }
    
    return [NSArray arrayWithArray:trustChain];
}


//获取服务器证书所有公钥
- (NSArray *)p_publicKeyTrustChainForServerTrust:(SecTrustRef)serverTrust
{
    NSMutableArray *publicKeys = [NSMutableArray array];
    
    // 获取证书链中的证书数量
    CFIndex certificateCount = SecTrustGetCertificateCount(serverTrust);
    
    for (CFIndex i = 0; i < certificateCount; i++) {
        // 获取证书链中的每个证书
        SecCertificateRef certificate = SecTrustGetCertificateAtIndex(serverTrust, i);
        
        // 提取公钥
        SecKeyRef publicKey = NULL;
        
        // iOS 10+ 可以直接从证书获取公钥
        if (@available(iOS 10.0, macOS 10.12, *)) {
            publicKey = SecCertificateCopyKey(certificate);
        } else {
            // iOS 10 以下需要创建临时信任对象来获取公钥
            SecTrustRef trust = NULL;
            SecPolicyRef policy = SecPolicyCreateBasicX509();
            
            if (SecTrustCreateWithCertificates(certificate, policy, &trust) == errSecSuccess) {
                // 评估信任对象
                SecTrustResultType result;
                if (SecTrustEvaluate(trust, &result) == errSecSuccess) {
                    // 从信任对象中提取公钥
                    publicKey = SecTrustCopyPublicKey(trust);
                }
            }
            
            if (trust) {
                CFRelease(trust);
            }
            
            if (policy) {
                CFRelease(policy);
            }
        }
        
        // 如果成功获取公钥，添加到数组
        if (publicKey) {
            [publicKeys addObject:(__bridge_transfer id)publicKey];
        }
    }
    return publicKeys;
}

- (BOOL)p_secKeyIsEqualToKey:(SecKeyRef)key1 key2:(SecKeyRef)key2
{
    return [(__bridge id)key1 isEqual:(__bridge id)key2];
}


#pragma mark - getter && setter
- (void)setPinnedCertificates:(NSSet<NSData *> *)pinnedCertificates
{
    _pinnedCertificates = pinnedCertificates;
    
    //根据证书数据获取对应的公钥
    if (pinnedCertificates) {
        NSMutableSet *mutablePinnedPublicKeys = [NSMutableSet setWithCapacity:pinnedCertificates.count];
        for (NSData *data in pinnedCertificates) {
            id publicKey = [self p_publicKeyFromCertificateData:data];
            if (publicKey) {
                [mutablePinnedPublicKeys addObject:publicKey];
            }
        }
        _pinnedPublicKeys = [mutablePinnedPublicKeys copy];
    }
    else {
        _pinnedPublicKeys = nil;
    }
}


@end
